Folder structure:¶
- Input(train) data 3422 images.
- Train csv : CSV containing information of bounding boxes around wheat spikes in each image.
- Test : Images on which prediction has to be made, prediction output has to be in the format present in submission csv.
- Submission csv : Format in which submission has to be made. It has 2 columns : 1. image name, 2. probability of bounding box containing wheat spike, correspnding co-orinates x,y,h,w ... prob and x,y,h,w for each predicted bbox.
- Intermediate folder (temporary folder to keep images to be processed batchwise).
Batchsize needs to be changed based on your system specifications.
Steps:¶
Step 1: In this kernal, I'll be using pyspark for data preparation as it requires manipulating 3000 images (train size). ~300 for kept for testing. Train csv does not contain information for all images present in Raw data, pls do not worry if numbers do not sum up.¶
For training, each image is converted in to binary image, where 0's represent background, 1's represent wheat spikes Encoding is done using bounding boxes coordinates present in the Train csv.
Step 2: Model (Keras vgg_unet) training on ~3000 images with n_classes=2 , input_height=1024, input_width=1024.¶
images are resized to reduced computation on CPU.
Step 3: Prediction on ~400 images to remove background (0) and only predict wheat spikes as foreground (1)¶
Step 4: Blob detection using blob_dog (difference of Gaussian) to detect blobs in model predicted segments for wheat spikes.¶
Step 5: IoU calculation, and computation of probability of each detection in each image.¶
Step 1: Data Preparation¶
Import libraries¶
In [1]:
import cv2
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
### Scikit learn libraries ###
from skimage.color import rgb2gray
from skimage.feature import blob_dog
from keras.models import model_from_json
### Pyspark libraries ###
import os,gc
import shutil
import pyspark
import itertools
from operator import add
from pyspark.sql.types import *
from pyspark import SparkContext
from pyspark.sql import DataFrame
from pyspark.sql.functions import udf
from pyspark.ml.image import ImageSchema
Reading input files¶
In [2]:
file = pd.read_csv("train_csv/train.csv") # This file contains bbox co-ordinates of wheat spikes for each image.
# Train csv has multiple rows for each image depending upon number of wheat spikes in each image.
ids = [file.iloc[i][0] for i in range(file.shape[0])]
# Factor by which we will resize each image as each raw image has size 1024*1024, which will be too large for network to train on to.
factor = 2; size = 512; batch_size = 500; x_axis_gaps = 5; coords_dist_th = 15; N =2999 ;N_ = N+30
Annotated samples¶
In [3]:
sample = os.listdir("train")[0:3]
for img_name in sample:
image = cv2.imread('train/'+ img_name)
# Get image ids for each image to access the segmentation co-ordinates
image_ids = [l for l,val in enumerate(ids) if val == str(img_name.split(".")[0])]
list_of_coords = [[int(float(val)) for val in file.iloc[v][3][1:len(file.iloc[v][3])-1].split(",")] for v in image_ids]
for l in list_of_coords:
image_annotated = cv2.rectangle(image,(int(l[0]),int(l[1])),(int((l[0]+l[2])),int((l[1]+l[3]))),(255,0,0),4)
fig, axes = plt.subplots(1, 2, figsize=(10, 5))#, sharex=True, sharey=True)
ax = axes.ravel()
image = cv2.imread('train/'+ img_name)
ax[0].imshow(cv2.cvtColor(image,cv2.COLOR_RGB2BGR))
ax[1].imshow(cv2.cvtColor(image_annotated,cv2.COLOR_RGB2BGR))
plt.show()
Let's get started with Data Preparation ................¶
Converting images to binary for training¶
- Data Preparation is the most critical part. For which, I had written a seperate module and imported the same here.
- Usually, our notebooks get consumed by "many Data scientists (referenced ones :)" , hence leaving data preparation part away to keep personal intellectual property safe.
In [4]:
if not os.path.exists('train_annotated'):
os.makedirs('train_annotated')
from data_prep_module import data_prep
Let's Spark it up!¶
Setting up Spark Context¶
Note, I made a few chnages in the backend code for Pyspark UDF to work here. In case you face challanges, feel free to comment/connect.
In [5]:
def run_spark():
sc=SparkContext(master="local[15]") # number[15] can be changed based on your system specification.
# print(sc.binaryFiles("spark_temp/*.png"))
image_df = sc.binaryFiles("spark_temp") # Reading images into Spark context
pyspark.sql.udf.UDFRegistration.register(name="data_prep", f = data_prep, returnType=StringType()) #registering UDF
# such that Spark context recongises the function used for data preparation.
job = [data_prep(ids,factor,file,size,x[0].split("/")[-1]) for idx,x in enumerate(image_df.take(batch_size))]
shutil.rmtree('spark_temp') #remove the folder once 'one batch' is complete to avoid Spark remembering
# indices for images it is done with.
sc.stop() # stop the spark context else Spark will have unnecessary information cached whcih we do not require anymore.
gc.collect() # remove any other cache which memory might have been holding.
In [6]:
input_list = os.listdir("train")[0:N]
input_seq = sorted(list(set([val for val in range(0, len(input_list),batch_size)]+ [len(input_list)])))
def image_transfer(value):
if not os.path.exists('spark_temp'):
os.makedirs('spark_temp')
image = cv2.imread('train/'+value)
cv2.imwrite("spark_temp/" + value.split(".")[0]+".jpg",cv2.resize(image,(size,size)))
for val in range(len(input_seq)-1):
img_batch = input_list[input_seq[val]:input_seq[val]+batch_size]
[image_transfer(val) for val in img_batch]
run_spark()
Data Preparation is Done!¶
Step 2: Model Training¶
Lets' do a bit checking on our data.
In [7]:
if not os.path.exists('train_temp'):
os.makedirs('train_temp')
def image_transfer_(value):
cv2.imwrite("train_temp/" + value.split(".")[0]+".png",cv2.resize(cv2.imread('train/'+value.split(".")[0]+".jpg"),(size,size)))
first_3000 = os.listdir("train_annotated")
task = [image_transfer_(val) for val in first_3000]
In [8]:
error_list = [val if np.max(cv2.imread('train_annotated/'+val)) > 1 else None for val in os.listdir("train_annotated")]
error_list = [val for val in error_list if val != None]
if len(error_list) > 0:
print(error_list)
[os.remove("train_temp/"+val.split(".")[0]+".jpg") for val in error_list]
[os.remove("train_annotated/"+val) for val in error_list]
In [9]:
# import tensorflow.compat.v1 as tensorflow
from keras_segmentation.models.unet import vgg_unet
model = vgg_unet(n_classes=2, input_height=size, input_width=size)
model.train(
train_images = "train_temp/",#train
train_annotations = "train_annotated/",n_classes = 2,epochs=10,steps_per_epoch=5,#annotations_prepped_train_v3
)
In [10]:
#saving model to disk
from keras.models import model_from_json
model_json = model.to_json()
with open("model.json", "w") as json_file:
json_file.write(model_json)
model.save_weights("model.h5")
In [11]:
for val in os.listdir("test/"):
cv2.imwrite("test/"+val.split(".")[0]+".jpg",cv2.resize(cv2.imread('test/'+val.split(".")[0]+".jpg"),(size,size)))
In [12]:
import warnings
warnings.simplefilter("ignore")
accuracy = []; submission = {"image_id" : [], "width" : [], "height" : [],"bbox" : [], "Prob":[]}
for idx,img_name in enumerate(os.listdir("test/")[0:5]):
if not os.path.exists('segmentation'):
os.makedirs('segmentation')
if not os.path.exists('Output'):
os.makedirs('Output')
out = model.predict_segmentation(
inp="test/"+ img_name,
out_fname="segmentation/" + img_name.split(".")[0] + ".png"
)
out_image = cv2.imread('segmentation/'+ img_name.split(".")[0] + ".png")
_image = cv2.imread('test/'+ img_name.split(".")[0] + ".jpg")
## Step 4:Blob detection
image_gray = rgb2gray(out_image)
blobs_dog = blob_dog(image_gray, max_sigma=30, threshold=.10)
try:
blobs_dog[:, 2] = blobs_dog[:, 2] * math.sqrt(2)
submission["image_id"].append(img_name.split(".")[0])
submission["width"].append(size); submission["height"].append(size)
submission["bbox"].append([str([int(blob[0]), int(blob[1]),int(blob[2]), int(blob[2])]) for blob in blobs_dog])
submission["Prob"].append(.50)
except:
submission["image_id"].append(img_name.split(".")[0])
submission["width"].append(size); submission["height"].append(size)
submission["bbox"].append([])
submission["Prob"].append(1)
blob_list_sorted = sorted([[int(blob[0]), int(blob[1]),int(blob[2])] for blob in blobs_dog])
len_val = len(blob_list_sorted)
for i,val in enumerate(range(0,len(blob_list_sorted))):
for j in range(i+1,len(blob_list_sorted)-2):
if max(i,j) < len_val and max([abs(np.diff(x)) for x in zip(blob_list_sorted[j][0:2], blob_list_sorted[i][0:2])]) < coords_dist_th:
blob_list_sorted = blob_list_sorted+[[min(x) for x in zip(blob_list_sorted[j], blob_list_sorted[i])][0:2]+[blob_list_sorted[j][2]+blob_list_sorted[i][2]]]
del blob_list_sorted[i]
del blob_list_sorted[j]
len_val = len(blob_list_sorted)
try:
for blob in blobs_dog:
y, x, r = blob
Iou = cv2.rectangle(_image,(int(x),int(y)),(int((x+r)),int((y+r))),(255,0,0),2)
# print(Iou)
cv2.imwrite("Output/" + img_name.split(".")[0] + ".png",Iou)
fig, axes = plt.subplots(1, 2, figsize=(12, 6))#, sharex=True, sharey=True)
ax = axes.ravel()
ax[0].imshow(cv2.cvtColor(cv2.imread('test/'+ img_name.split(".")[0] + ".jpg"),cv2.COLOR_RGB2BGR))
ax[1].imshow(Iou)
plt.show()
except:pass
pd.DataFrame(submission).to_csv("submission.csv")
In [13]:
pd.DataFrame(submission).to_csv("submission.csv")
Results are looking alryt! Let's test on few more images..
In [14]:
if not os.path.exists('train_50'):
os.makedirs('train_50')
for val in os.listdir("train/")[N:N_]:
cv2.imwrite("train_50/"+val.split(".")[0]+".png",cv2.resize(cv2.imread('train/'+val.split(".")[0]+".jpg"),(size,size)))
In [15]:
warnings.simplefilter("ignore")
for idx,img_name in enumerate(os.listdir("train_50")):
if not os.path.exists('segmentation'):
os.makedirs('segmentation')
out = model.predict_segmentation(
inp="train_50/"+ img_name,
out_fname="segmentation/" + img_name.split(".")[0] + ".png"
)
out_image = cv2.imread('segmentation/'+ img_name.split(".")[0] + ".png")
_image = cv2.imread('train_50/'+ img_name.split(".")[0] + ".png")
## Step 4:Blob detection
image_gray = rgb2gray(out_image)
blobs_dog = blob_dog(image_gray, max_sigma=30, threshold=.10)
blob_list_sorted = sorted([[int(blob[0]), int(blob[1]),int(blob[2])] for blob in blobs_dog])
len_val = len(blob_list_sorted)
for i,val in enumerate(range(0,len(blob_list_sorted))):
for j in range(i+1,len(blob_list_sorted)-2):
if max(i,j) < len_val and max([abs(np.diff(x)) for x in zip(blob_list_sorted[j][0:2], blob_list_sorted[i][0:2])]) < coords_dist_th:
blob_list_sorted = blob_list_sorted+[[min(x) for x in zip(blob_list_sorted[j], blob_list_sorted[i])][0:2]+[blob_list_sorted[j][2]+blob_list_sorted[i][2]]]
del blob_list_sorted[i]
del blob_list_sorted[j]
len_val = len(blob_list_sorted)
try:
for blob in blobs_dog:
y, x, r = blob
Iou = cv2.rectangle(_image,(int(x),int(y)),(int((x+r)),int((y+r))),(255,0,0),2)
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
ax = axes.ravel()
ax[0].imshow(cv2.cvtColor(cv2.imread('train_50/'+ img_name.split(".")[0] + ".png"),cv2.COLOR_RGB2BGR))
ax[1].imshow(Iou)
plt.show()
except:pass
End Notes¶
- There are classes in image dataset, few images with too much sunlight, few images with clear distinction between wheat heads and background.
- It will be required for me to train the model with additional sophistication to be able to predict wheat heads in different light conditions, different orientation due to wind and color variation depending on the age of the crop.
- Look forward to next notebook on the same..
To all my fellow Data Scientists, non privileged ones, privileged ones, in the industry because of relatives, or by yourself... Happy coding ;)
This work took me a month time to identify right pre-processing technique which worked for all (almost) images. if you notice, results are pretty satisfactory in terms of algorithm picking all the wheat heads in an image despite the varied size of the wheat heads, different light conditions and orientation.
I agree, there is scope for improvement (which I will take up later in sometime..) but I am equally proud that the solution is able to seperate wheat heads from leaves which are of same color and at times, same shape and size. and top that solution is able to discard the background which is soil, dry leaves and etc..
Few Applications of this exercise:
- Health care image analysis in detecting specific type cell (cancerous vs non cancerous). https://www.robots.ox.ac.uk/~vgg/research/cell_detection/images/histo.png
- Head count in an image.
- Automated vechicle counting in a parking lot.
- Automated inventory counting in a store (inventory of one type).
- And...If you are really interested, counting stars in the sky :D
Appreciate the effort.., do critic, share feedback. You all are most welcome.
In [ ]: